# 安装 pip install -U matplotlib
import matplotlib.pyplot as plt
# 安装 pip install networkx
import networkx as nx
# networkx的初步使用 https://blog.csdn.net/qq_45956730/article/details/127352658


def demo_1():
    g = nx.DiGraph()  # 初始化一个有向图
    # 向图添加边，每条边由两个节点和一个权重数据组成
    g.add_weighted_edges_from([['a', 'b', 5], ['a', 'c', 4], ['b', 'a', 5], ['b', 'd', 3], ['d', 'c', 1]])
    # 画图
    nx.draw(g, with_labels=True)
    # 生成的图表保存到本地
    plt.savefig('./img/demo_1.png')
    # 预览
    plt.show()
    # 关闭画布
    plt.close()


def demo_2():
    # 添加画布
    fig, ax = plt.subplots(1, 1)

    # 生成图并初始化图的数据
    g = nx.DiGraph()
    g.add_weighted_edges_from([['a', 'b', 5], ['a', 'c', 4], ['b', 'a', 5], ['b', 'd', 3], ['d', 'c', 1]])

    # 生成节点坐标信息，用于画图
    pos = nx.spring_layout(g, seed=5)
    # 生成边的标签数据，用于绘制边标签
    edge_labels = {(d[0], d[1]): d[2]['weight'] for d in g.edges(data=True)}

    # 绘制图表，图表有一个整体的draw和四部分组成
    nx.draw(g, pos, node_size=[500, 400, 300, 200], node_color='#8888ff', width=5, edge_color='#8888ff')
    nx.draw_networkx_edges(g, pos, width=1, edge_color='#0088ff')
    nx.draw_networkx_edge_labels(g, pos, edge_labels=edge_labels, font_size=20)
    nx.draw_networkx_nodes(g, pos, node_size=150)
    nx.draw_networkx_labels(g, pos, font_size=20)

    # 给图设置一个标题，并保存
    ax.set_title('networkx test')
    # 生成的图表保存到本地
    plt.savefig('./img/demo_2.png')
    # 预览图表
    plt.show()
    # 关闭画布
    plt.close()


# 绘制带权无向图
def demo_3():

    # 问题 2：无向图的最短路问题（司守奎，数学建模算法与应用，P43，例4.3）
    G2 = nx.Graph()  # 创建无向图
    G2.add_weighted_edges_from([(1, 2, 2), (1, 3, 8), (1, 4, 1),
                                (2, 3, 6), (2, 5, 1),
                                (3, 4, 7), (3, 5, 5), (3, 6, 1), (3, 7, 2),
                                (4, 7, 9),
                                (5, 6, 3), (5, 8, 2), (5, 9, 9),
                                (6, 7, 4), (6, 9, 6),
                                (7, 9, 3), (7, 10, 1),
                                (8, 9, 7), (8, 11, 9),
                                (9, 10, 1), (9, 11, 2),
                                (10, 11, 4)])  # 向图中添加多条赋权边: (node1,node2,weight)

    # 两个指定顶点之间的最短加权路径
    minWPath_v1_v4 = nx.dijkstra_path(G2, source=1, target=5)  # 顶点 0 到 顶点 3 的最短加权路径
    print("顶点 v1 到 顶点 v5 的最短加权路径: ", minWPath_v1_v4)
    # 两个指定顶点之间的最短加权路径的长度
    lminWPath_v1_v4 = nx.dijkstra_path_length(G2, source=1, target=5)  # 最短加权路径长度
    print("顶点 v1 到 顶点 v5 的最短加权路径长度: ", lminWPath_v1_v4)
    pos = nx.spring_layout(G2)  # 用 FR算法排列节点
    nx.draw(G2, pos, with_labels=True, alpha=0.5)
    labels = nx.get_edge_attributes(G2, 'weight')
    nx.draw_networkx_edge_labels(G2, pos, edge_labels=labels)
    plt.savefig('./img/demo_3.png')
    plt.show()
    plt.close()


# 绘制带权有向图
def demo_4():
    import matplotlib.pyplot as plt  # 导入 Matplotlib 工具包
    import networkx as nx  # 导入 NetworkX 工具包

    # 问题 2：无向图的最短路问题（司守奎，数学建模算法与应用，P43，例4.3）
    G2 = nx.DiGraph()  # 创建：空的 有向图

    G2.add_edge(1, 2, weight=1)  # 添加 带权边，weight表示边权
    G2.add_edge(5, 3, weight=7)
    G2.add_edge(2, 3, weight=4)
    G2.add_edge(3, 4, weight=3)
    G2.add_edge(7, 9, weight=4)
    G2.add_edge(3, 5, weight=5)
    G2.add_edge(4, 7, weight=9)

    # 两个指定顶点之间的最短加权路径
    minWPath_v1_v5 = nx.dijkstra_path(G2, source=1, target=5)  # 顶点 0 到 顶点 3 的最短加权路径
    print("顶点 v1 到 顶点 v5 的最短加权路径: ", minWPath_v1_v5)
    # 两个指定顶点之间的最短加权路径的长度
    lMinWPath_v1_v5 = nx.dijkstra_path_length(G2, source=1, target=5)  # 最短加权路径长度
    print("顶点 v1 到 顶点 v5 的最短加权路径长度: ", lMinWPath_v1_v5)
    pos = nx.spring_layout(G2)  # 用 FR算法排列节点
    nx.draw(G2, pos, with_labels=True, alpha=0.5)
    labels = nx.get_edge_attributes(G2, 'weight')
    nx.draw_networkx_edge_labels(G2, pos, edge_labels=labels)
    plt.savefig('./img/demo_4.png')
    plt.show()
    plt.close()


def demo_5():
    # 问题 2：无向图的最短路问题（司守奎，数学建模算法与应用，P43，例4.3）
    G2 = nx.DiGraph()  # 创建：空的 有向图

    G2.add_edge(1, 2, weight=1)  # 添加 带权边，weight表示边权
    G2.add_edge(5, 3, weight=7)
    G2.add_edge(2, 3, weight=4)
    G2.add_edge(3, 4, weight=3)
    G2.add_edge(7, 9, weight=4)
    G2.add_edge(3, 5, weight=5)
    G2.add_edge(4, 7, weight=9)

    # 两个指定顶点之间的最短加权路径
    minWPath_v1_v5 = nx.dijkstra_path(G2, source=1, target=5)  # 顶点 0 到 顶点 3 的最短加权路径
    print("顶点 v1 到 顶点 v5 的最短加权路径: ", minWPath_v1_v5)
    # 两个指定顶点之间的最短加权路径的长度
    lMinWPath_v1_v5 = nx.dijkstra_path_length(G2, source=1, target=5)  # 最短加权路径长度
    print("顶点 v1 到 顶点 v5 的最短加权路径长度: ", lMinWPath_v1_v5)
    pos = nx.spring_layout(G2)  # 用 FR算法排列节点
    nx.draw(G2, pos, with_labels=True, alpha=0.5)
    labels = nx.get_edge_attributes(G2, 'weight')
    nx.draw_networkx_edge_labels(G2, pos, edge_labels=labels)
    # plt.show()

    edgeList = []
    for i in range(len(minWPath_v1_v5) - 1):
        edgeList.append((minWPath_v1_v5[i], minWPath_v1_v5[i + 1]))
    nx.draw_networkx_edges(G2, pos, edgelist=edgeList, edge_color='m', width=4)  # 设置边的颜色
    plt.savefig('./img/demo_5.png')
    plt.show()
    plt.close()


if __name__ == '__main__':
    # demo_1()
    # demo_2()
    demo_3()
    # demo_4()
    # demo_5()
